#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
@author: xjx
@time: 2024/4/30 10:36 
@file: load_analysis_3.py
@project: nanchangproject
@describe: 使用神经网络进行载荷预测
"""
import torch
import torch.nn as nn


class NeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(NeuralNetwork, self).__init__()
        # 定义隐藏层和输出层的线性变换
        self.hidden = nn.Linear(input_size, hidden_size)  # 输入到隐藏层的线性变换
        self.output = nn.Linear(hidden_size, output_size)  # 隐藏层到输出层的线性变换
        # 定义激活函数
        self.sigmoid = nn.Sigmoid()  # 使用 Sigmoid 激活函数

    def forward(self, X):
        # 正向传播过程
        hidden_output = self.sigmoid(self.hidden(X))  # 隐藏层的输出，使用 Sigmoid 激活函数
        output = self.sigmoid(self.output(hidden_output))  # 输出层的输出，使用 Sigmoid 激活函数
        return output


# 定义预测函数
def predict(model, new_data, device):
    prediction = model(new_data.to(device))
    return prediction.item()


class LoadAnalysisBPModel(object):
    # 用于神经网络模型的调用
    def __init__(self):
        self.input_size = 4
        self.hidden_size = 5
        self.output_size = 1
        self.device = torch.device('cpu')
        self.model = NeuralNetwork(self.input_size, self.hidden_size, self.output_size)
        self.model.load_state_dict(torch.load('trained_model.pth', map_location=torch.device('cpu')))
        self.speed_max = 97.4375  # 最大速度
        self.speed_min = 0  # 最小速度
        self.actual_torque_max = 100  # 最大净输出
        self.actual_torque_min = 0  # 最小净输出
        self.friction_torque_max = 16  # 最大摩擦扭矩
        self.friction_torque_min = 5  # 最小摩擦扭矩
        self.engine_speed_max = 3524.954545  # 最大发动机转速
        self.engine_speed_min = 514.5454545  # 最小发动机转速
        self.load_max = 6115  # 最大载荷
        self.load_min = 3036  # 最小载荷
        self.speed_range = self.speed_max - self.speed_min
        self.actual_torque_range = self.actual_torque_max - self.actual_torque_min
        self.friction_torque_range = self.friction_torque_max - self.friction_torque_min
        self.engine_speed_range = self.engine_speed_max - self.engine_speed_min
        self.load_range = self.load_max - self.load_min

    @staticmethod
    def change(num):
        if num > 1:
            return 1
        elif num < 0:
            return 0
        else:
            return num

    def normalization(self, speed: float, actual_torque: float, friction_torque: float, engine_speed: float):
        """
        用于归一化数据
        @param speed:
        @param actual_torque:
        @param friction_torque:
        @param engine_speed:
        @return:
        """
        normalized_speed = LoadAnalysisBPModel.change((speed - self.speed_min) / self.speed_range)
        normalized_actual_torque = LoadAnalysisBPModel.change(
            (actual_torque - self.actual_torque_min) / self.actual_torque_range)
        normalized_friction_torque = LoadAnalysisBPModel.change(
            (friction_torque - self.friction_torque_min) / self.friction_torque_range)
        normalized_engine_speed = LoadAnalysisBPModel.change(
            (engine_speed - self.engine_speed_min) / self.engine_speed_range)
        return normalized_speed, normalized_actual_torque, normalized_friction_torque, normalized_engine_speed

    def analysis(self, speed: float, actual_torque: float, friction_torque: float, engine_speed: float):
        normalized_speed, normalized_actual_torque, normalized_friction_torque, normalized_engine_speed = \
            self.normalization(speed, actual_torque, friction_torque, engine_speed)
        data = torch.tensor(
            [[normalized_speed, normalized_actual_torque, normalized_friction_torque, normalized_engine_speed]],
            dtype=torch.float32).to(self.device)
        prediction = self.model(data).item()  # 预测结果
        result = prediction * self.load_range + self.load_min  # 反归一化后的结果
        return result


if __name__ == "__main__":
    # 初始化神经网络
    # input_size = 4
    # hidden_size = 5
    # output_size = 1
    # my_nn = NeuralNetwork(input_size, hidden_size, output_size)
    # my_nn.load_state_dict(torch.load('trained_model.pth', map_location=torch.device('cpu')))
    # device = torch.device('cpu')
    # my_nn = my_nn.to(device)
    # new_data = torch.tensor([[1, 1, 1, 1]], dtype=torch.float32).to(device)  # 输入x向量
    # output = predict(my_nn, new_data, device)
    # print("预测结果:", output)
    a = LoadAnalysisBPModel()
    result = a.analysis(speed=0, actual_torque=8, friction_torque=5, engine_speed=800.5)
    print(f'预测结果:{round(result,4)}千克')
